import os
import jieba
import pandas as pd
import joblib
import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from keras.models import load_model
from keras.preprocessing.sequence import pad_sequences
from keras.utils import to_categorical


# 定义数据目录
PREDICTION_FOLDER = 'predict_data/'  # 新的文件夹用于保存待预测的文本数据

# 中文分词：使用jieba进行分词
def chinese_tokenizer(text):
    return list(jieba.cut(text))

# 加载训练好的模型和工具
lr_model = joblib.load('lr_model.pkl')  # Logistic Regression模型
lstm_model = load_model('lstm_model_best.keras')  # LSTM模型
vectorizer = joblib.load('tfidf_vectorizer.pkl')  # TF-IDF 向量化器
label_encoder = joblib.load('label_encoder.pkl')  # 标签编码器


# 预测函数：根据输入文本返回标签
def predict_text(texts, model_type='lr'):
    # 文本向量化
    x_data_vectorized = vectorizer.transform(texts).toarray()
    
    # 对于LSTM模型，填充序列，使每个文本长度一致
    max_length = 100  # 假设文本最大长度为100
    x_data_padded = pad_sequences(x_data_vectorized, padding='post', maxlen=max_length)

    # 根据指定的模型类型选择预测方式
    if model_type == 'lr':
        # 使用Logistic Regression进行预测
        predictions = lr_model.predict(x_data_vectorized)
    elif model_type == 'lstm':
        # 使用LSTM进行预测
        predictions = lstm_model.predict(x_data_padded)
        predictions = np.argmax(predictions, axis=1)  # 取概率最大的类
    
    # 将预测的标签解码
    predicted_labels = label_encoder.inverse_transform(predictions)
    return predicted_labels

# 处理文件夹中的所有待预测的文本文件
for filename in os.listdir(PREDICTION_FOLDER):
    if filename.endswith('.xlsx'):
        file_path = os.path.join(PREDICTION_FOLDER, filename)
        df = pd.read_excel(file_path)  # 读取待预测的文本数据
        
        # 提取文本数据
        texts_to_predict = df['text'].values.tolist()

        # 进行预测
        predicted_labels = predict_text(texts_to_predict, model_type='lr')  # 使用Logistic Regression模型

        # 将预测结果加入原数据
        df['predicted_label'] = predicted_labels

        # 保存带有预测标签的文件到新的文件夹
        df.to_excel(os.path.join(PREDICTION_FOLDER, f'predicted_{filename}'), index=False)

print("所有预测已完成，结果已保存到预测文件夹中。")
